import requests
import threading
import time
import logging
from typing import Callable, Dict, Any, Optional

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class AgentSDK:
    """智能体SDK核心类"""
    
    def __init__(self, platform_url: str, agent_config: Dict[str, Any]):
        """
        初始化SDK
        
        Args:
            platform_url: 平台地址，如 http://localhost:8080
            agent_config: 智能体配置
                - name: 智能体名称
                - type: 智能体类型
                - description: 描述（可选）
                - capabilities: 能力列表（可选）
                - endpoint: 访问端点
        """
        self.platform_url = platform_url.rstrip('/')
        self.agent_config = agent_config
        self.agent_id: Optional[str] = None
        self.token: Optional[str] = None
        self.heartbeat_thread: Optional[threading.Thread] = None
        self.running = False
        self.task_callback: Optional[Callable] = None
        
    def register(self) -> Dict[str, str]:
        """注册智能体到平台"""
        url = f"{self.platform_url}/api/agent/register"
        data = {
            "agentName": self.agent_config["name"],
            "agentType": self.agent_config["type"],
            "description": self.agent_config.get("description", ""),
            "capabilities": str(self.agent_config.get("capabilities", [])),
            "endpoint": self.agent_config["endpoint"]
        }
        
        try:
            response = requests.post(url, json=data, timeout=10)
            response.raise_for_status()
            result = response.json()
            
            if result["code"] == 200:
                self.agent_id = result["data"]["agentId"]
                self.token = result["data"]["token"]
                logger.info(f"智能体注册成功: {self.agent_id}")
                return result["data"]
            else:
                raise Exception(f"注册失败: {result['message']}")
        except Exception as e:
            logger.error(f"注册失败: {e}")
            raise
    
    def start_heartbeat(self, interval: int = 30):
        """
        启动心跳
        
        Args:
            interval: 心跳间隔（秒），默认30秒
        """
        if not self.agent_id:
            raise Exception("请先注册智能体")
        
        self.running = True
        self.heartbeat_thread = threading.Thread(
            target=self._heartbeat_loop,
            args=(interval,),
            daemon=True
        )
        self.heartbeat_thread.start()
        logger.info("心跳已启动")
    
    def _heartbeat_loop(self, interval: int):
        """心跳循环"""
        url = f"{self.platform_url}/api/agent/heartbeat"
        
        while self.running:
            try:
                data = {
                    "agentId": self.agent_id,
                    "status": "ONLINE"
                }
                response = requests.post(url, json=data, timeout=10)
                response.raise_for_status()
                logger.debug(f"心跳发送成功: {self.agent_id}")
            except Exception as e:
                logger.error(f"心跳发送失败: {e}")
            
            time.sleep(interval)
    
    def receive_task(self, callback: Callable[[Dict[str, Any]], None]):
        """
        设置任务接收回调函数
        
        Args:
            callback: 任务处理回调函数，接收任务数据作为参数
        """
        self.task_callback = callback
        logger.info("任务回调已设置")
    
    def submit_result(self, task_id: str, status: str, output_data: Any = None, error_message: str = None):
        """
        提交任务结果
        
        Args:
            task_id: 任务ID
            status: 任务状态 (SUCCESS/FAILED)
            output_data: 输出数据
            error_message: 错误信息（如果失败）
        """
        url = f"{self.platform_url}/api/task/callback"
        data = {
            "taskId": task_id,
            "status": status,
            "outputData": str(output_data) if output_data else None,
            "errorMessage": error_message
        }
        
        try:
            response = requests.post(url, json=data, timeout=10)
            response.raise_for_status()
            logger.info(f"任务结果提交成功: {task_id}")
        except Exception as e:
            logger.error(f"任务结果提交失败: {e}")
            raise
    
    def unregister(self):
        """注销智能体"""
        if not self.agent_id:
            return
        
        self.running = False
        if self.heartbeat_thread:
            self.heartbeat_thread.join(timeout=5)
        
        url = f"{self.platform_url}/api/agent/unregister"
        try:
            response = requests.post(url, params={"agentId": self.agent_id}, timeout=10)
            response.raise_for_status()
            logger.info(f"智能体注销成功: {self.agent_id}")
        except Exception as e:
            logger.error(f"智能体注销失败: {e}")
    
    def __enter__(self):
        """上下文管理器入口"""
        self.register()
        self.start_heartbeat()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器退出"""
        self.unregister()
